# -*- coding: utf-8 -*-

"""
Created on Thu Apr 29 10:08:51 2021

@author: shijie
"""

import numpy as np
import os
from multiprocessing import Pool
import time
from Bio.PDB.PDBParser import PDBParser
import multiprocessing as mp
import matplotlib.pyplot as plt
import warnings
from Bio.PDB.NeighborSearch import NeighborSearch as NG
from Bio.PDB import Superimposer as Super

warnings.filterwarnings("ignore")

def register_parameters():
    
    from argparse import ArgumentParser
    
    parser=ArgumentParser()
    
    
    parser.add_argument('-data_path', '--data_path'
                        ,default='/fsa/home/ww_dengcw/Desktop/RNAS/new_sample1/'
                        ,type=str, help='path of pdb files ')
    
    parser.add_argument('-index_path', '--index_path'
                        ,default='/fsa/home/ww_dengcw/Desktop/INDEX_FILES/'
                        ,type=str, help='output: path of index files ')

    parser.add_argument('-path_workdir', '--path_workdir'
                        ,default='/fsa/home/ww_dengcw/Desktop/GNN_multiprocessing_distrubuted/'
                        ,type=str, help='path of source code')
    
    
    parser.add_argument('-radius', '--radius'
                        ,default=14
                        ,type=float, help='radius of local enviroment. initial value is 14A. you had better not to change it.')
    
    parser.add_argument('-n', '--n'
                        ,default=4
                        ,type=int, help='cpu cores to compute index files.')
    
    parser.add_argument('-nc', '--nc'
                        ,default=2
                        ,type=int, help='cpu cores to check and re-compute index files.')
    
    args=parser.parse_args()
    
    
    global path_data, path_index_files, path_workdir, path_rmss, path_rmsjpg
    
    # paths
    path_data=args.data_path
    path_index_files=args.index_path
    path_workdir=args.path_workdir
    path_rmss=path_workdir+'all_rmsd.npy'
    path_rmsjpg=path_workdir+'all_rmsd_distribution.jpg'
    
    # parameters
    global radius, nprocess, rmsd_range, check_nprocess
    radius= args.radius
    nprocess=args.n
    check_nprocess=args.nc
    
    print('----------register paths and parameters--------')
    print('data path: ', path_data)
    print('-----------------------------')
    print('indexs path: ', path_index_files)
    print('-----------------------------')
    print('src path: ', path_workdir)
    print('-----------------------------')   
    print('local radius: ', radius)
    print('-----------------------------')  
    print('number of processes: ', nprocess)
    print('-----------------------------')
    print('number of check processes: ', check_nprocess)
    print('-----------------------------') 
    
register_parameters()

def get_res(s):
    if 'A' in s:
        return 'A'
    elif 'G' in s:
        return 'G'
    elif 'C' in s:
        return 'C'
    elif 'U' in s:
        return 'U'
    else:
        return 'A'
    
def tab(s, id1, max_gap=4):
    s=iter(s)
    out=''
    for i in range(len(id1)-1):
        gap=id1[i+1]-id1[i]
        out+=next(s)
        if abs(gap)<=max_gap:          
            for j in range(gap-1):
                out+='N'
    out+=next(s)
    return out
    
def c_i(s1, s2):
    a=0
    for i in range(len(s1)):
        if (s1[i]==s2[i]):
            a+=1
    return a

def get_env_index(s1, s2):
    l1=len(s1)
    l2=len(s2)
    
    c=abs(l1-l2)
    s=0
    env1, env2=[], []
    if l1<l2:     
        cis=[]
        for s in range(c):
            ci=c_i(s2[s:s+l1], s1)
            cis.append(ci)
            if ci==1.:
                break
        s=cis.index(max(cis))
        
        p1=0
        p2=len(s2[:s].replace('N', ''))
        sp1=0
        sp2=s          
    elif l1>l2:
        cis=[]
        for s in range(c):
            ci=c_i(s1[s:s+l2], s2)
            cis.append(ci)
            if ci==1.:
                break
        s=cis.index(max(cis))
        
        p1=len(s1[:s].replace('N', ''))
        p2=0
        sp1=s
        sp2=0
    else:
        p1=0
        p2=0
        sp1=0
        sp2=0
    for i in range(l1 if l1<l2 else l2):
        if s1[sp1+i]==s2[sp2+i] and s1[sp1+i]!='N':
            env1.append(p1)
            env2.append(p2)
        if s1[sp1+i]!='N':
            p1+=1
        if s2[sp2+i]!='N':
            p2+=1 
    return env1, env2
        
def align(stand_ress, ress, env=None):
    if not env:
        if len(stand_ress)==len(ress):
            return stand_ress, ress
        else:
            s1, s2='', ''
            id1, id2=[], []
            for r in stand_ress:
                s1+=get_res(r.resname.strip())
                id1.append(r.id[1])
            for r in ress:
                s2+=get_res(r.resname.strip())
                id2.append(r.id[1])
            s11=tab(s1, id1)
            s22=tab(s2, id2)
            env1, env2=get_env_index(s11, s22)
            Res1=[stand_ress[i] for i in env1]
            Res2=[ress[j] for j in env2]
    else:
        if len(stand_ress)==len(ress):
            Res1=[ress[i] for i in env]
            Res2=[stand_ress[i] for i in env]
        else:
            s1, s2='', ''
            id1, id2=[], []
            for r in stand_ress:
                s1+=get_res(r.resname.strip())
                id1.append(r.id[1])
            for r in ress:
                s2+=get_res(r.resname.strip())
                id2.append(r.id[1])
            s11=tab(s1, id1)
            s22=tab(s2, id2)
            env1, env2=get_env_index(s11, s22)
            Res1=[stand_ress[i] for i in env1]
            Res2=[ress[j] for j in env2]           
            id3=[r.id[1] for r in Res2]
            
            id2=[id2[i] for i in env]
            idset=list(set(id3).intersection(set(id2)))
            env=[id3.index(x) for x in idset]
            
            Res1=[Res1[i] for i in env]
            Res2=[Res2[j] for j in env]            
    return Res1, Res2

def search_index(names, name):
    try:
        return names.index(name)
    except:
        return names.index(name[0]+name[2]+name[1])

def RMS(ress1,ress2): 

    atoms1=[]
    atoms2=[]
    for i in range(len(ress1)):
        ra1=list(ress1[i].get_atoms())
        name1=[a.name.strip() for a in ra1]
        ra2=list(ress2[i].get_atoms())
        name2=[a.name.strip() for a in ra2]
        
        setname=list(set(name1)&set(name2))
        temp1, temp2=[], []
        for a in ra1:
            if a.name.strip() in setname:
                temp1.append(a)
        for a in ra2:
            if a.name.strip() in setname:
                temp2.append(a)
        temp1.sort(key=lambda x:setname.index(x.name.strip()))
        temp2.sort(key=lambda x:setname.index(x.name.strip()))
        atoms1.extend(temp1)
        atoms2.extend(temp2)  
    sup=Super()
    sup.set_atoms(atoms1,atoms2)
    return sup.rms

def get_model_rms(rna_dir, path_sample):
    rna_model=rna_dir.split(sep="/")[-1]
    rna=rna_model[:4]
    model_index=0 if len(rna_model)<9 else int(rna_model[6:].replace(".pdb", ""))
    
    if not model_index:
        model_index=0
        rmsd=0
    else:
        parser=PDBParser()
        
        origin=list(parser.get_structure("0", path_sample+rna+"/"+rna+".pdb").get_residues())
        struct=list(parser.get_structure("0", rna_dir).get_residues())
        Res1, Res2=align(origin, struct)
        model_index=int(rna_model[6:].replace(".pdb", "")) 
        rmsd=RMS(Res1, Res2)
    print(os.getpid(), rna, model_index, rmsd)
    return (rna, model_index, rmsd)

def px(a):
    return a[0]

def absorb(a):
    b=a.copy()
    for [x,y] in a:
        b.append([y,x])
    b.sort(key=px)
    center=b[0][0]
    env=[center]
    Micro_envs=[]
    for [x,y] in b:
        if x==center:
            env.append(y)
        else:
            Micro_envs.append(env)
            center=x
            env=[center,y]
    Micro_envs.append(env)
    return Micro_envs

def rename(rna,m_name):
    rna=rna.replace(".pdb","")
    m_name=m_name.replace(".pdb","")
    if m_name[0]=="S":
        m_name="s"+str((int(m_name[2:])+1))
    elif rna==m_name:
        m_name="s0"
    else:
        m_name=m_name[5:]
    return rna,m_name

def res_to_atoms(ress):
    #残基列表转原子列表
    out=[]
    for res in ress:
        out.extend(list(res.get_atoms()))
    return out

def to_str(env):
    out=""
    for x in env:
        out+=str(x)+" "
    return out.strip()
        
def get_envs(rna,m_name,stand_ress,rmss,radius,D1,D2, all_env=False):
    parser=PDBParser()
    model=parser.get_structure("1", D1+rna+"/"+m_name)
    ress=list(model.get_residues())
    ng=NG(list(model.get_atoms()),10)
    n=len(ress)
    
    pair_res=ng.search_all(radius,"R")
    pair_num=[[ress.index(a),ress.index(b)] for (a,b) in pair_res]
    Micro_envs=absorb(pair_num)
    Micro_envs.sort(key=len,reverse=True)
    
    if all_env:
        out1=Micro_envs
    else:
    #取nres max的部分
        out=[]
        out1=[]
        for env in Micro_envs:
            out1.append(env)
            out.extend(env)
            if len(set(out))==n:
                break   
    out=""
    #此处可用多线程？？？
    rna,m_name=rename(rna,m_name)
    #id2=[x.id[1] for x in stand_ress]
    for i,env in enumerate(out1):
        if stand_ress:
            Res1, Res2=align(stand_ress, ress, env)
            rms=RMS(Res1,Res2)
        else:
            rms=0
        rmss.append(rms)
        if i!=len(out1):
            out+=str(rms)+" "+to_str(env)+"\n"
        else:
            out+=str(rms)+" "+to_str(env)
    try:
        f=open(D2+rna+"/"+rna+"_"+m_name+".txt","w")
    except:
        os.mkdir(D2+rna+"/")
        f=open(D2+rna+"/"+rna+"_"+m_name+".txt","w")
    f.write(out)
    f.close()
    print(rna,m_name,"---当前进程:",os.getpid())
    
def get_num_of_rnamodel():
    out={}
    RNAS=os.listdir(path_data)
    for rna in RNAS:
        models=os.listdir(path_data+rna+"/")
        cnt=0
        for m in models:
            if '.pdb' in m:
                cnt+=1
        out[rna]=cnt
    return out

def main_getenvname(n_process):
    rms_range=100
    try:
        os.mkdir(path_index_files)
    except:
        pass
    RNAS=os.listdir(path_data)
    old_RNAS=os.listdir(path_index_files)
    new=[]
    for rna in RNAS:
        if rna not in old_RNAS:
            new.append(rna)
    if not new:
        print("has no new RNA, exit.")
        return
    print("new: ", len(new))
    pool=Pool(n_process)
    parser=PDBParser()
    m=mp.Manager()
    rmss=m.list()
    for rna in new[:]:
        model_names=os.listdir(path_data+rna+"/")
        stand_ress=list((parser.get_structure("1",path_data+rna+"/"+rna+".pdb")).get_residues())
        for m_name in model_names:
            if len(m_name)!=8:
                pool.apply_async(func=get_envs,args=(rna,m_name,stand_ress,rmss,radius,path_data, path_index_files, True))
            else:
                pool.apply_async(func=get_envs,args=(rna,m_name,None,rmss,radius,path_data, path_index_files, True))
    pool.close()
    pool.join()
    all_rmss=np.array(rmss)
    np.save(path_rmss, all_rmss)
    
    fig=plt.figure(dpi=1000, figsize=(10,6))
    ax=fig.add_subplot(1,1,1)
    ax.hist(all_rmss, bins=rms_range, range=(0,rms_range), density=True, color="r", edgecolor="black")
    ax.set_xlabel("RMSD A")
    ax.set_title("sample rmsd distrubution")
    plt.savefig(path_rmsjpg, dpi=1000)
    print("all RNAS has been completed")
    
def check_data(n_process):
    num_models=get_num_of_rnamodel()
    
    print("Now checking data integrity--------------")
    RNA_ENVS=os.listdir(path_index_files)

    pool=Pool(n_process)
    parser=PDBParser()
    m=mp.Manager()
    rmss=m.list()
    
    re_RNAS=[]
    for rna in RNA_ENVS:
        ENVS=os.listdir(path_index_files+rna+"/")
        if len(ENVS)!=num_models[rna]:
            re_RNAS.append(rna)
            print(rna, "has just %d/%d envs, Now recreating it"%(len(ENVS), num_models[rna]))
            os.system("rm -rf "+path_index_files+rna+"/")
    if not re_RNAS:
        print("all RNAS are integrity, exit.")
        return 1
    for rna in re_RNAS:
        model_names=os.listdir(path_data+rna+"/")
        stand_ress=list((parser.get_structure("1",path_data+rna+"/"+rna+".pdb")).get_residues())
        for m_name in model_names:
            if len(m_name)<=9:
                stand_ress=None
            pool.apply_async(func=get_envs,args=(rna,m_name,stand_ress,rmss,radius,path_data, path_index_files, True))
            
    pool.close()
    pool.join()
    return 0

def test_time(n_proccess):
    start=time.time()
    main_getenvname(n_proccess)
    end=time.time()
    inf=str(n_proccess)+" CPU spend "+str(end-start)+" Seconds"
    print(inf)
    print(end-start,"Seconds")
    """
    f=open("/home/dcheng/test_time.txt","w")
    f.write(sskk)
    f.close()
    """
    
if __name__=="__main__":

    main_getenvname(nprocess)
    
    ### check error
    r=check_data(check_nprocess)   
    cn=0
    while r!=1:
        r=check_data(check_nprocess)
        cn+=1
        if cn>=3:
            break
    print('----completed--------')